from openai import OpenAI
import base64
import requests
from dotenv import load_dotenv
import os

load_dotenv()

def image_path_to_base64(image_path):
    with open(image_path, "rb") as image_file:
        return base64.b64encode(image_file.read()).decode('utf-8')

# 使用示例（替换为你的实际路径）
local_image_path = "your_image.jpg"  # 本地图像路径
base64_image = image_path_to_base64(local_image_path)

client = OpenAI(
    # 需替换为实际的 API 基础 URL
    base_url="https://ai.gitee.com/v1",
    api_key=os.getenv("api_key"),
)

response = client.images.generate(
    extra_body={
        "image": base64_image,  # 使用转换后的base64字符串
        "height": 1024,
        "steps": 25,
        "guidance_scale": 6,
    },
    prompt="generate a same style image",
    model="Kolors",
    size="1024x1024",
)

for i, image_data in enumerate(response.data):
    if image_data.url:
        ext = image_data.url.split('.')[-1].split('?')[0] or "jpg"
        filename = f"Kolors-output-{i}.{ext}"
        response_download = requests.get(image_data.url, timeout=30)
        response_download.raise_for_status()
        with open(filename, "wb") as f:
            f.write(response_download.content)
        print(f"Downloaded image to {filename}")
    elif image_data.b64_json:
        image_bytes = base64.b64decode(image_data.b64_json)
        filename = f"Kolors-output-{i}.jpg"
        with open(filename, "wb") as f:
            f.write(image_bytes)
        print(f"Saved image to {filename}")